(*
   Copyright (C) 2024 International Digital Economy Academy.
   This program is licensed under the MoonBit Public Source
   License as published by the International Digital Economy Academy,
   either version 1 of the License, or (at your option) any later
   version. This program is distributed in the hope that it will be
   useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the MoonBit
   Public Source License for more details. You should have received a
   copy of the MoonBit Public Source License along with this program. If
   not, see
   <https://www.moonbitlang.com/licenses/moonbit-public-source-license-v1>.
*)

(* 错误格式化模块，提供更友好的错误报告 *)

module Io = Basic_io
module Lst = Basic_lst
module Vec = Basic_vec

(* ANSI颜色代码 *)
let use_color = ref true

let red s = if !use_color then "\027[31m" ^ s ^ "\027[0m" else s
let green s = if !use_color then "\027[32m" ^ s ^ "\027[0m" else s
let yellow s = if !use_color then "\027[33m" ^ s ^ "\027[0m" else s
let blue s = if !use_color then "\027[34m" ^ s ^ "\027[0m" else s
let magenta s = if !use_color then "\027[35m" ^ s ^ "\027[0m" else s
let cyan s = if !use_color then "\027[36m" ^ s ^ "\027[0m" else s
let bold s = if !use_color then "\027[1m" ^ s ^ "\027[0m" else s

(* 禁用颜色 *)
let disable_color () = use_color := false

(* 启用颜色 *)
let enable_color () = use_color := true

(* 获取文件内容的行 *)
let get_file_lines file =
  try
    let content = Io.load_file file in
    String.split_on_char '\n' content
  with _ -> []

(* 格式化位置信息 *)
let format_location loc =
  let open Loc in
  let file = loc.loc_start.pos_fname in
  let line = loc.loc_start.pos_lnum in
  let start_col = loc.loc_start.pos_cnum - loc.loc_start.pos_bol in
  let end_col = loc.loc_end.pos_cnum - loc.loc_start.pos_bol in
  (file, line, start_col, end_col)

(* 生成错误指示器 *)
let generate_indicator start_col end_col =
  let spaces = String.make start_col ' ' in
  let length = max 1 (end_col - start_col) in
  let carets = String.make length '^' in
  spaces ^ red carets

(* 格式化错误消息 *)
let format_error err =
  match err with
  | Local_diagnostics.Unbound_value (loc, id) ->
      let (file, line, start_col, end_col) = format_location loc in
      let lines = get_file_lines file in
      let line_content = try List.nth lines (line - 1) with _ -> "" in
      let indicator = generate_indicator start_col end_col in
      let message = Printf.sprintf "%s:%d:%d: %s: Unbound value '%s'\n%s\n%s\n%s"
        file line start_col
        (bold (red "Error"))
        (Longident.to_string id)
        line_content
        indicator
        (yellow "Hint: Check if the variable name is spelled correctly or if it needs to be imported.")
      in
      message

  | Local_diagnostics.Type_error (loc, expected, actual) ->
      let (file, line, start_col, end_col) = format_location loc in
      let lines = get_file_lines file in
      let line_content = try List.nth lines (line - 1) with _ -> "" in
      let indicator = generate_indicator start_col end_col in
      let message = Printf.sprintf "%s:%d:%d: %s: Type mismatch\n%s\n%s\n%s\n%s"
        file line start_col
        (bold (red "Error"))
        line_content
        indicator
        (cyan (Printf.sprintf "Expected: %s" (Stype.string_of_type expected)))
        (cyan (Printf.sprintf "Actual: %s" (Stype.string_of_type actual)))
      in
      message

  | Local_diagnostics.Syntax_error (loc, msg) ->
      let (file, line, start_col, end_col) = format_location loc in
      let lines = get_file_lines file in
      let line_content = try List.nth lines (line - 1) with _ -> "" in
      let indicator = generate_indicator start_col end_col in
      let message = Printf.sprintf "%s:%d:%d: %s: %s\n%s\n%s"
        file line start_col
        (bold (red "Error"))
        msg
        line_content
        indicator
      in
      message

  | _ ->
      (* 默认情况下使用原始错误格式化 *)
      Local_diagnostics.string_of_error err

(* 格式化警告消息 *)
let format_warning warn =
  match warn with
  | Warnings.Unused_variable (loc, name) ->
      let (file, line, start_col, end_col) = format_location loc in
      let lines = get_file_lines file in
      let line_content = try List.nth lines (line - 1) with _ -> "" in
      let indicator = generate_indicator start_col end_col in
      let message = Printf.sprintf "%s:%d:%d: %s: Unused variable '%s'\n%s\n%s\n%s"
        file line start_col
        (bold (yellow "Warning"))
        name
        line_content
        indicator
        (blue "Hint: Consider removing this variable or prefix it with '_' to indicate it's intentionally unused.")
      in
      message

  | _ ->
      (* 默认情况下使用原始警告格式化 *)
      Warnings.string_of_warning warn

(* 打印错误 *)
let print_error err =
  let formatted = format_error err in
  Printf.eprintf "%s\n\n" formatted

(* 打印警告 *)
let print_warning warn =
  let formatted = format_warning warn in
  Printf.eprintf "%s\n\n" formatted

(* 提供修复建议 *)
let suggest_fix err =
  match err with
  | Local_diagnostics.Unbound_value (_, id) ->
      Some (Printf.sprintf "Did you mean to import '%s' or declare it first?" (Longident.to_string id))
  
  | Local_diagnostics.Type_error (_, expected, actual) ->
      if Stype.is_int expected && Stype.is_float actual then
        Some "Try using 'Int.from_float' to convert from Float to Int."
      else if Stype.is_float expected && Stype.is_int actual then
        Some "Try using 'Float.from_int' to convert from Int to Float."
      else if Stype.is_string expected then
        Some "Try using '.to_string()' to convert to String."
      else
        None
  
  | _ -> None

(* 打印诊断信息 *)
let print_diagnostics diagnostics =
  let errors = Diagnostics.get_errors diagnostics in
  let warnings = Diagnostics.get_warnings diagnostics in
  
  if not (Vec.is_empty errors) then begin
    Printf.eprintf "%s\n" (bold (red (Printf.sprintf "Found %d error(s):" (Vec.length errors))));
    Vec.iter (fun err -> print_error err) errors;
  end;
  
  if not (Vec.is_empty warnings) then begin
    Printf.eprintf "%s\n" (bold (yellow (Printf.sprintf "Found %d warning(s):" (Vec.length warnings))));
    Vec.iter (fun warn -> print_warning warn) warnings;
  end;
  
  (* 提供修复建议 *)
  if not (Vec.is_empty errors) then begin
    let suggestions = Vec.filter_map suggest_fix errors in
    if not (Vec.is_empty suggestions) then begin
      Printf.eprintf "%s\n" (bold "Suggestions:");
      Vec.iter (fun suggestion -> Printf.eprintf "  - %s\n" suggestion) suggestions;
    end;
  end